# 1. 分析dwd.csv文件，提取城市、明星、最低价、最高价等信息
# 2. 预测下次演唱会的最低价和最高价
# 3. 输出结果到predict_result.csv文件
# predict_result.csv我用来绘制折线图，要体现出预测价格以及历史价格的变化趋势# ...existing code...
import re
from datetime import timedelta
import numpy as np
import pandas as pd

INPUT_CSV = "../dwd.csv"
OUTPUT_CSV = "../predict_result.csv"

def parse_price(x):
    try:
        return int(str(x).split('-')[0])
    except:
        return np.nan

def parse_price_max(x):
    try:
        return int(str(x).split('-')[-1])
    except:
        return np.nan

def extract_start_date(time_str):
    if pd.isna(time_str):
        return pd.NaT
    # 常见格式：2025.05.02-05.04 或 2025.04.19 周六 14:00 或 2025.05.10-05.11
    m = re.search(r'(\d{4})[\./-](\d{1,2})[\./-](\d{1,2})', str(time_str))
    if m:
        y, mo, d = m.groups()
        try:
            return pd.to_datetime(f"{y}-{int(mo):02d}-{int(d):02d}")
        except:
            return pd.NaT
    return pd.NaT

def predict_for_group(gdf):
    # 输入：一个分组的 DataFrame（已按 date 排序），返回带预测行的 DataFrame
    df = gdf.sort_values("start_date").copy()
    # 如果没有解析到日期或仅有一条记录，使用平均作为预测
    n = len(df)
    if n == 0:
        return df
    if n == 1 or df['start_date'].isna().all():
        pred_min = int(round(df['min'].iloc[-1]))
        pred_max = int(round(df['max'].iloc[-1]))
        pred_date = pd.NaT
    else:
        # 时间序列回归：用时间戳对min和max做线性拟合
        valid = df.dropna(subset=['start_date'])
        if len(valid) >= 2:
            ts = valid['start_date'].astype('int64')  # 纳秒
            # 预测下一时间点：最近时间 + 中位间隔
            diffs = valid['start_date'].diff().dt.days.dropna()
            median_interval = int(diffs.median()) if len(diffs)>0 else 30
            next_date = valid['start_date'].max() + pd.Timedelta(days=median_interval)
            x = ts.values.astype(np.float64)
            x0 = next_date.value
            # 拟合 min
            y_min = valid['min'].values.astype(np.float64)
            coeff_min = np.polyfit(x, y_min, 1)
            pred_min_f = np.polyval(coeff_min, x0)
            # 拟合 max
            y_max = valid['max'].values.astype(np.float64)
            coeff_max = np.polyfit(x, y_max, 1)
            pred_max_f = np.polyval(coeff_max, x0)
            pred_min = int(round(max(pred_min_f, 0)))
            pred_max = int(round(max(pred_max_f, pred_min)))
            pred_date = next_date
        else:
            # 虽然有多个记录但无完整时间信息，退回到均值预测
            pred_min = int(round(df['min'].mean()))
            pred_max = int(round(df['max'].mean()))
            pred_date = pd.NaT

    pred_row = {
        'city': df['city'].iloc[0],
        'users': df['users'].iloc[0],
        'title': df['title'].iloc[0],
        'start_date': pred_date,
        'time': 'PREDICTED' if pd.isna(pred_date) else pred_date.strftime('%Y-%m-%d'),
        'min': pred_min,
        'max': pred_max,
        'is_predicted': True
    }
    df['is_predicted'] = False
    return pd.concat([df, pd.DataFrame([pred_row])], ignore_index=True, sort=False)

def main():
    df = pd.read_csv(INPUT_CSV, dtype=str)
    # 基本列清洗
    df['min'] = df['min'].apply(lambda x: pd.to_numeric(x, errors='coerce'))
    df['max'] = df['max'].apply(lambda x: pd.to_numeric(x, errors='coerce'))
    df['start_date'] = df['time'].apply(extract_start_date)
    # 标准化明星列（users），空值填为"未知"
    df['users'] = df['users'].fillna('未知').replace('', '未知')
    # 保留需要列
    df2 = df[['city', 'users', 'title', 'start_date', 'time', 'min', 'max']].copy()

    results = []
    # 以 city + users 分组预测
    for (city, users), g in df2.groupby(['city', 'users']):
        res = predict_for_group(g)
        results.append(res)

    all_res = pd.concat(results, ignore_index=True, sort=False)
    # 将 start_date 转为字符串便于绘图/保存
    all_res['start_date_str'] = all_res['start_date'].dt.strftime('%Y-%m-%d')
    # 对预测行如果无日期，用 time 字段标注
    all_res.loc[all_res['is_predicted'] & all_res['start_date'].isna(), 'start_date_str'] = all_res.loc[all_res['is_predicted'] & all_res['start_date'].isna(), 'time']

    # 排序便于查看：按 city, users, start_date（NaT 放后）
    all_res['start_date_sort'] = all_res['start_date'].apply(lambda x: x.value if pd.notna(x) else 9e18)
    all_res = all_res.sort_values(['city', 'users', 'start_date_sort']).drop(columns=['start_date_sort'])

    # 输出 CSV（包含历史和预测）
    all_res.to_csv(OUTPUT_CSV, index=False, columns=['city','users','title','start_date_str','time','min','max','is_predicted'])
    print(f"已生成预测结果：{OUTPUT_CSV}")

if __name__ == "__main__":
    main()
# ...existing code...